#include <math.h>
#include <iostream>
#include <NvOnnxParser.h>
#include "../include/misc.h"

using namespace std;

void print_ioinfo(nvinfer1::INetworkDefinition *network);

void print_ioinfo(nvinfer1::ICudaEngine *engine);

float cos_similarity(const float* A, const float* B,const int& len);

int main() {
    srand(0);
    string root_path = "/files/code/Csdn-Source-Code-Implementation/tensorrt/step0-diymodel/weights/";
    string trt_file = root_path + "best2.trt";

    TRTLogger logger;

    std::vector<uint8_t> data = load_file(trt_file);

    nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
    nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(data.data(), data.size());
    if (engine == nullptr) {
        printf("deserialize failed\n");
        runtime->destroy();
        return -1;
    }
    print_ioinfo(engine);

    nvinfer1::Dims dim = engine->getBindingDimensions(0);
    int length = dim.d[1];
    int channel = dim.d[2];
    int input_capacity = 1 * length * channel;
    float *input_host = new float[input_capacity];
    float *input_device = nullptr;
    cudaMalloc(reinterpret_cast<void **>(&input_device), input_capacity * sizeof(float));

    nvinfer1::Dims dim1 = engine->getBindingDimensions(1);
    int feats_capacity = 1 * dim1.d[1];
    vector<float> feats_host(feats_capacity, -1);
    float *feats_device = nullptr;
    cudaMalloc(reinterpret_cast<void **>(&feats_device), feats_capacity * sizeof(float));

    nvinfer1::Dims dim2 = engine->getBindingDimensions(2);
    int output_capacity = 1 * dim2.d[1];
    float output_host[output_capacity];
    float *output_device = nullptr;
    cudaMalloc(reinterpret_cast<void **>(&output_device), output_capacity * sizeof(float));

    vector<float *> binding{input_device, feats_device, output_device};

    // input is 10 feats, 5 channel in one feat
    for (int i = 0; i < length; i++) {
        for (int j = 0; j < channel; j++) {
            *(input_host + i * channel + j) = float(std::rand()) / float(RAND_MAX);
        }
    }
    vector<float> weight{1, 0.5, 1, -1, 1};
    float bias = 0;
    float *mid_val = new float[10];
    float *output = new float(-100000);
    for (int i = 0; i < length; i++) {
        *(mid_val + i) = 0.0;
        for (int j = 0; j < channel; j++) {
            *(mid_val + i) += *(input_host + i * channel + j) * weight[j];
        }
        *(mid_val + i) += bias;
        *output = max(*output, *(mid_val + i));
    }

    cudaStream_t stream = nullptr;
    cudaStreamCreate(&stream);
    cudaMemcpyAsync(input_device, input_host, input_capacity * sizeof(float), cudaMemcpyHostToDevice, stream);

    nvinfer1::IExecutionContext *context = engine->createExecutionContext();

    // 明确当前推理时，使用的数据输入大小
    auto input_dims = context->getBindingDimensions(0);
    input_dims.d[0] = 1;
    // 设置当前推理时，input大小
    context->setBindingDimensions(0, input_dims);
    bool success = context->enqueueV2((void **) binding.data(), stream, nullptr);

    cudaMemcpyAsync(output_host, output_device, output_capacity * sizeof(float), cudaMemcpyDeviceToHost, stream);

    cudaStreamSynchronize(stream);

    *output_host = exp(*output_host);

    float simularity_a = cos_similarity(output_host, output, output_capacity);

    printf("simularity_a: %f\n",simularity_a);

    context->destroy();
    engine->destroy();
    runtime->destroy();
    printf("done\n");
    return 0;
}

void print_ioinfo(nvinfer1::INetworkDefinition *network) {
    int input_nums = network->getNbInputs();
    for (int i = 0; i < input_nums; i++) {
        nvinfer1::ITensor *input = network->getInput(i);
        string input_name = input->getName();
        nvinfer1::Dims input_dims = input->getDimensions();
        printf("input<%d> tensor name = %s, dim = ", i, input_name.c_str());
        for (int i = 0; i < input_dims.nbDims; i++) {
            cout << input_dims.d[i] << "  ";
        }
        cout << endl;
    }
    int output_nums = network->getNbOutputs();
    for (int i = 0; i < output_nums; i++) {
        nvinfer1::ITensor *output = network->getOutput(i);
        string output_name = output->getName();
        nvinfer1::Dims output_dims = output->getDimensions();
        printf("output<%d> tensor name = %s,  dim = ", i, output_name.c_str());
        for (int i = 0; i < output_dims.nbDims; i++) {
            cout << output_dims.d[i] << "  ";
        }
        cout << endl;
    }
}

void print_ioinfo(nvinfer1::ICudaEngine *engine) {
    std::cout << "infos of engine:" << std::endl;
    std::cout << "bind nums:" << engine->getNbBindings() << std::endl;
    for (int i = 0; i < engine->getNbBindings(); i++) {
        std::cout << engine->getBindingName(i) << " dims: ";
        nvinfer1::Dims dim = engine->getBindingDimensions(i);
        for (int i = 0; i < dim.nbDims; i++) {
            cout << dim.d[i] << "  ";
        }
        std::cout << std::endl;
    }
}

float cos_similarity(const float* A, const float* B,const int& len){
    float A_scale=0.0;
    float B_scale=0.0;
    float AB=0.0;
    for(int i=0;i<len;i++){
        A_scale += *(A+i) * *(A+i);
        B_scale += *(B+i) * *(B+i);
        AB += *(A+i) * *(B+i);
    }
    float cos = AB / (sqrt(A_scale)*sqrt(B_scale));
    return cos;
}
